import pandas as pd
from sklearn.metrics import mean_absolute_error
from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.model_selection import GridSearchCV
import joblib  # 用于保存模型

# 读取数据
data = pd.read_csv('data.csv')

# 数据预处理（转换时间、排序、创建滞后特征、删除NaN）
data['time'] = pd.to_datetime(data['time'])
data = data.sort_values(by='time').reset_index(drop=True)
data['IC_DCS_CHIMNEY_NOX_lagged'] = data['IC_DCS_CHIMNEY_NOX'].shift(3)
data = data.dropna()

# 分离特征和标签
y = data['IC_DCS_CHIMNEY_NOX_lagged']
X = data.drop(columns=['IC_DCS_CHIMNEY_NOX', 'IC_DCS_CHIMNEY_NOX_lagged', 'time'])

# 基于时间戳划分训练集和测试集
# 我们将使用一个固定的时间点来划分，而不是随机划分
# 假设我们使用最后20%的时间段作为测试集
train_cutoff = data['time'].quantile(0.8)
X_train = X[data['time'] < train_cutoff]
y_train = y[data['time'] < train_cutoff]
X_test = X[data['time'] >= train_cutoff]
y_test = y[data['time'] >= train_cutoff]

# 初始化HistGradientBoostingRegressor（设置一些基础参数，但将在网格搜索中调优）
model_base = HistGradientBoostingRegressor(max_iter=100, learning_rate=0.1, max_leaf_nodes=31, random_state=42)

# 定义参数网格（根据HistGradientBoostingRegressor的参数进行调整）
param_grid = {
    'max_iter': [100, 200, 300],  # 迭代次数（树的生长轮数）
    'learning_rate': [0.01, 0.1, 0.2],  # 学习率
    'max_leaf_nodes': [31, 100, 200],  # 最大叶子节点数
    'min_samples_leaf': [1, 2, 5],  # 每个叶子节点的最小样本数
    'l2_regularization': [0.0, 0.1, 0.5]  # L2正则化强度
}

# 注意：HistGradientBoostingRegressor不支持某些传统GBM的参数，如max_depth和min_samples_split。
# 相反，它使用max_leaf_nodes和min_samples_leaf来控制树的复杂度。

# 初始化GridSearchCV
grid_search = GridSearchCV(estimator=model_base, param_grid=param_grid, scoring='neg_mean_absolute_error', cv=3, n_jobs=-1)

# 训练模型（进行网格搜索）
grid_search.fit(X_train, y_train)

# 获取最佳参数和模型
best_params = grid_search.best_params_
best_model = grid_search.best_estimator_  # 使用GridSearchCV找到的最佳模型

# 预测测试集
y_pred = best_model.predict(X_test)

# 保存模型
joblib.dump(best_model, "xgboost_model.joblib")

# 计算MAE
mae = mean_absolute_error(y_test, y_pred)
print(f'Mean Absolute Error: {mae}')